import numpy as np
import matplotlib.pyplot as plt

def gso(base):
    print base.shape
    new_base=np.zeros_like((base), dtype=float)
    new_base[0]=base[0]
    for i in range(1, len(base)):
        new_base[i]=base[i]
        for k in range(i):
            new_base[i]=new_base[i]-proj(new_base[k], base[i])
        new_base[i]=new_base[i]/np.sqrt(np.dot(new_base[i], new_base[i]))
    new_base[0]=new_base[0]/np.sqrt(np.dot(new_base[0], new_base[0]))
    
    return new_base

def proj(u, v):
    pro=np.dot(v, u)/np.dot(u, u)*u
    return pro

class cbc():
    def __init__(self, chainLen=8, truncate=8):
        self.chain_length = chainLen
        self.hamiltonian = [2.0*np.identity(self.chain_length) \
            - np.diag(np.ones(self.chain_length-1),1) \
            - np.diag(np.ones(self.chain_length-1), -1) for i in range(4)]
        self.hamiltonian[1][-1,-1] = 1
        self.hamiltonian[2][0,0] = 1
        self.hamiltonian[3][[[0,-1],[0,-1]]] = 1

        self.site_num = chainLen

        self.tunelling = np.zeros( (self.chain_length, self.chain_length) )
        self.tunelling[-1, 0] = -1.0

        self.m = truncate        
        
        self.diagonalize()
        self.RG_energy = []
        self.RG_state = []
        for i in range(4):
            for j in range(self.m/4):
                self.RG_energy.append(self.energy[i][j])
                self.RG_state.append(self.state[i][:, j])
        self.RG_energy=np.array(self.RG_energy)
        self.RG_state=gso(np.array(self.RG_state).conj())
        
        self.wavefunction = self.RG_state

    def diagonalize(self):
        self.energy = []
        self.state = []
        for hamil in self.hamiltonian:
            [u,v]=np.linalg.eigh(hamil)
            idx = u.argsort()
            self.energy.append(u[idx])
            self.state.append(v[:, idx])

    def block_renormalize(self):
        self.RG_hamiltonian = [np.dot(self.RG_state, \
            np.dot(hamil, self.RG_state.T.conj())) for hamil in self.hamiltonian]
        self.RG_tunelling = np.dot( self.RG_state, \
            np.dot(self.tunelling, self.RG_state.T.conj()) )

        self.block_grow()

        self.diagonalize()
        self.RG_energy = []
        self.RG_state = []
        for i in range(4):
            for j in range(self.m/4):
                self.RG_energy.append(self.energy[i][j])
                self.RG_state.append(self.state[i][:, j])
        self.RG_energy=np.array(self.RG_energy)
        self.RG_state=gso(np.array(self.RG_state).conj())
        self.wavefunction = np.dot(self.wavefunction_dbl.T, self.RG_state.T).T

    def block_grow(self):
        self.hamiltonian=[np.r_[np.c_[self.RG_hamiltonian[0],self.RG_tunelling],
                                np.c_[self.RG_tunelling.conj().T, self.RG_hamiltonian[0]]],
                          np.r_[np.c_[self.RG_hamiltonian[0],self.RG_tunelling],
                                np.c_[self.RG_tunelling.conj().T, self.RG_hamiltonian[1]]],
                          np.r_[np.c_[self.RG_hamiltonian[2],self.RG_tunelling],
                                np.c_[self.RG_tunelling.conj().T, self.RG_hamiltonian[0]]],
                          np.r_[np.c_[self.RG_hamiltonian[2],self.RG_tunelling],
                                np.c_[self.RG_tunelling.conj().T, self.RG_hamiltonian[1]]]]
        zeroMat= np.zeros((self.m, self.m))
        self.tunelling = \
            np.r_[ \
                np.c_[zeroMat, zeroMat],\
                np.c_[self.RG_tunelling, zeroMat] ]
        
        zeroWave = np.zeros(self.site_num)
        wave1 = [np.array([x, zeroWave]).flatten() for x in self.wavefunction]
        wave2 = [np.array([zeroWave, x]).flatten() for x in self.wavefunction]
        self.wavefunction_dbl= (np.r_[wave1, wave2])
        
        self.site_num = self.site_num * 2
